import argparse
import torch
import cv2
import numpy as np
import onnxruntime as ort
from numpy import random
from models.experimental import attempt_load
from utils.general import check_img_size, check_requirements, check_imshow, non_max_suppression, scale_coords
from utils.torch_utils import select_device
from utils.datasets import LoadStreams, LoadImages
from utils.plots import plot_one_box

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--onnx', nargs='+', type=str, default='/home/zhicheng.luo/Weights/chuwei_best_20240805.onnx',
                        help='model.pt path(s)')
    parser.add_argument('--weights', nargs='+', type=str, default='/home/zhicheng.luo/Weights/chuwei_best_20240805.pt',
                        help='model.pt path(s)')
    parser.add_argument('--source', type=str, default='/home/zhicheng.luo/data/chuwei/toilet_44.jpg',
                        help='source')  # file/folder, 0 for webcam
    parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
    parser.add_argument('--conf-thres', type=float, default=0.35, help='object confidence threshold')
    parser.add_argument('--iou-thres', type=float, default=0.65, help='IOU threshold for NMS')
    parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    parser.add_argument('--view-img', action='store_true', default=True, help='display results')
    parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 0 2 3')
    parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
    parser.add_argument('--augment', action='store_true', help='augmented inference')
    opt = parser.parse_args()

    device = select_device(opt.device)
    model = attempt_load(opt.weights, map_location=device)  # load FP32 model
    stride = int(model.stride.max())  # model stride
    imgsz = check_img_size(opt.img_size, s=stride)  # check img_size
    dataset = LoadImages(opt.source, img_size=imgsz, stride=stride)

    # Get names and colors
    names = model.module.names if hasattr(model, 'module') else model.names
    colors = [[random.randint(0, 255) for _ in range(3)] for _ in names]

    #onnx
    cuda = True
    # providers = ['CPUExecutionProvider']
    # providers = ['CUDAExecutionProvider']
    # providers = ['TensorrtExecutionProvider', ('CUDAExecutionProvider', {'cudnn_conv_algo_search': 'DEFAULT',})]
    # providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] if cuda else ['CPUExecutionProvider']
    providers = [('CUDAExecutionProvider', {'cudnn_conv_algo_search': 'DEFAULT', }),
                 'CPUExecutionProvider'] if cuda else ['CPUExecutionProvider']  # EXHAUSTIVE#HEURISTIC#DEFAULT
    # providers = ['TensorrtExecutionProvider', 'CUDAExecutionProvider']
    session = ort.InferenceSession(opt.onnx, providers=providers)

    for path, img, im0s, vid_cap in dataset:
        img = img.astype(np.float32)
        img /= 255.0  # 0 - 255 to 0.0 - 1.0
        img = torch.from_numpy(img)
        if img.ndimension() == 3:
            img = img.unsqueeze(0)
            img = img.numpy()

        outname = [i.name for i in session.get_outputs()]  # ['output']
        inname = [i.name for i in session.get_inputs()]
        inp = {inname[0]: img}
        outputs = session.run(outname, inp)[0]
        print('over')